#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# @Description: 压测
import re

from jinja2 import Environment, FileSystemLoader

import constants
from modules.operationBase import OperationBase
from utils.logger import Logger

log = Logger()


def prepare_jinja2_conf(template_file, data):
    env = Environment(loader=FileSystemLoader(searchpath=constants.templates_path))
    template = env.get_template(template_file)
    rendered_config = template.render(data)
    log.success(f"模版制作完成：\n{rendered_config}")
    return rendered_config


def build_bench_capacity(capacity: str, percent: float) -> str:
    """
    传入当前容量和百分比，返回以G为单位的整数容量字符串。

    :param capacity: 当前容量，例如 "16T" 或 "500G"
    :param percent: 百分比，例如 25 表示 25%
    :return: 以G为单位的整数容量字符串，例如 "400G"
    """
    # 使用正则表达式提取数值和单位
    match = re.match(r"^(\d+\.?\d*)([TtGg]?)$", capacity)
    if not match:
        raise ValueError("Invalid capacity format")

    value = float(match.group(1))  # 提取数值部分
    unit = match.group(2).upper()  # 提取单位部分并转换为大写

    # 根据单位转换为G
    if unit == 'T':
        value_in_gb = value * 1024  # 1T = 1024G
    elif unit == 'G' or unit == '':  # 如果没有单位，默认为G
        value_in_gb = value  # 已经是G单位
    else:
        raise ValueError("Unsupported unit. Only 'T' and 'G' are supported.")

    # 计算百分比后的容量
    result_gb = value_in_gb * (percent / 100)

    # 转换为整数并返回字符串
    return f"{int(result_gb)}G"


class QBench(OperationBase):
    def __init__(self, parse_args):
        super().__init__(parse_args)

    def bench_fio_big_file(self, percent):
        """
        fio制作大文件
        1、先检查当前的qdfs文件系统有多少容量
        2、根据容量，指定文件，25%、50%、75%、90%
        :param percent: 25、50、75%、90
        :return: 返回压测的大小
        """
        # 获取挂载点可用容量
        capacity = self.ssh.exec_cmd(constants.SYSTEM_DF_TH_CAPACITY.format(self.point))
        log.info(f"当前挂载点{self.point}可用容量为：{capacity}")
        size = build_bench_capacity(capacity.strip(''), percent)

        data = {
            'mount_point': self.point,
            'fio_big_size': size
        }
        rendered_config = prepare_jinja2_conf(constants.FIO_BIG_FILE_TEMPLATE, data)
        self.ssh.upload_template_file(rendered_config, constants.FIO_BIG_FILE)

        self.ssh.exec_cmd(command=constants.BENCH_FIO.format(constants.FIO_BIG_FILE))
        return size

    def bench_vdbench_small_file(self):
        """
        vdbench制作小文件
        :return:
        """
        data = {
            'mount_point': self.point
        }
        rendered_config = prepare_jinja2_conf(constants.VDBENCH_SMALL_FILE_TEMPLATE, data)
        self.ssh.upload_template_file(rendered_config, constants.VDBENCH_SMALL_FILE)

        self.ssh.exec_cmd(command=constants.BENCH_VDBENCH_FILE.format(constants.VDBENCH_SMALL_FILE))

    def bench_fallocate_file(self, percent):
        """
        使用fallocate创建一个文件
        :param percent: 25、50、75%、90
        :return: 返回size大小
        """
        capacity = self.ssh.exec_cmd(constants.SYSTEM_DF_TH_CAPACITY.format(self.point))
        log.info(f"当前{self.point}，可用容量为{capacity}")
        size = build_bench_capacity(capacity, percent)
        log.info(f"fallocate size is {size}")
        self.ssh.exec_cmd(command=constants.BENCH_FALLOCATE.format(size, self.point))
        return size

